import numpy as np


# np.random.seed(1)


def viterbi(mat):
    # 取前两个
    hidden_status_num, period = mat.shape
    positions = np.zeros((hidden_status_num, period))
    D1 = mat[:, 0]
    distance = np.zeros((hidden_status_num, period))
    positions[:, 0] = [i for i in range(len(D1))]
    distance[:, 0] = D1
    D2 = mat[:, 1]
    for i, val in enumerate(D2):
        acc_current = val + D1
        positions[i, 1] = np.argmin(acc_current)
        distance[i, 1] = np.min(acc_current)

    for p in range(2, period):
        D = mat[:, p]
        for i, val in enumerate(D):
            acc_current = val + distance[:, p-1]
            positions[i, p] = np.argmin(acc_current)
            distance[i, p] = np.min(acc_current)

    print("positions:", positions)
    print("distance:", distance)

    min_dis_idx = np.argmin(distance[:, -1])
    indices = [0 for i in range(period)]
    for p in range(period - 1, -1, -1):
        if p == period - 1:
            indices[p] = min_dis_idx
        else:
            indices[p] = int(positions[indices[p+1], p+1])
    str_indices = [str(i) for i in indices]
    route = "->".join(str_indices)
    return route


hidden_status_num = 3
period = 5


mat = np.random.randint(0, high=100, size=(hidden_status_num, period))
print(mat)
route = viterbi(mat)
print(route)


